#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Side-Channel Attack基Class
定义Side-Channel Attack的统一接口
"""

import numpy as np
import time
import logging
from abc import ABC, abstractmethod
from typing import Dict, Any, Optional, List, Tuple
from dataclasses import dataclass


@dataclass
class AttackResult:
    """AttackResultDataClass"""
    attack_type: str
    target_byte: int
    execution_time: float
    success: bool
    confidence: float
    best_hypothesis: Optional[int] = None
    top_candidates: Optional[List[Tuple[int, float]]] = None  # [(hypothesis, score), ...]
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}
        if self.top_candidates is None:
            self.top_candidates = []
    
    def get_rank_of_hypothesis(self, hypothesis: int) -> Optional[int]:
        """
        获取指定False设In候选List中的排名
        
        Args:
            hypothesis: 要Find的False设值
            
        Returns:
            排名（From1Start），IfNotInList中ReturnNone
        """
        for rank, (hyp, score) in enumerate(self.top_candidates, start=1):
            if hyp == hypothesis:
                return rank
        return None
    
    def is_hypothesis_in_top_n(self, hypothesis: int, n: int = 10) -> bool:
        """
        Check指定False设YesNoIn前n个候选中
        
        Args:
            hypothesis: 要Check的False设值
            n: 前n个候选
            
        Returns:
            YesNoIn前n个候选中
        """
        rank = self.get_rank_of_hypothesis(hypothesis)
        return rank is not None and rank <= n


class BaseSideChannelAttack(ABC):
    """Side-Channel Attack基Class"""
    
    def __init__(self, name: str):
        """
        InitializeAttack器
        
        Args:
            name: Attack器名称
        """
        self.name = name
        self.logger = logging.getLogger(f"SCA.{name}")
        self.traces = None
        self.plaintexts = None
        self.ciphertexts = None
        self.keys = None
        self.timestamps = None
        self.durations = None
        
    def load_data(self, npz_file: str) -> bool:
        """
        LoadDataFile
        
        Args:
            npz_file: NPZDataFilePath
            
        Returns:
            YesNoLoadSuccess
        """
        try:
            self.logger.info(f"Loading data from: {npz_file}")
            data = np.load(npz_file, allow_pickle=True)
            
            # ExtractData
            self.plaintexts = data['plaintexts']
            self.ciphertexts = data.get('ciphertexts')
            self.keys = data.get('keys')  # Side-Channel Attack时KeyYesOptional的
            
            # ProcessPowerTraceData - Support多种Format
            if 'power_traces' in data:
                self.traces = data['power_traces']
            elif 'traces' in data:
                self.traces = data['traces']
            elif 'cpu_usage' in data:
                # ProcessCPUUse率Data作ForPowerTrace
                cpu_data = data['cpu_usage']
                
                # 找To最大长度并填充AllTraceTo相同长度
                max_len = max(len(trace) for trace in cpu_data)
                padded_traces = []
                
                for trace in cpu_data:
                    # 将ListConvertFornumpyArray并填充To最大长度
                    trace_array = np.array(trace, dtype=np.float64)
                    if len(trace_array) < max_len:
                        # 用Last一个值填充
                        padded = np.pad(trace_array, (0, max_len - len(trace_array)), 
                                      mode='edge')
                    else:
                        padded = trace_array
                    padded_traces.append(padded)
                
                self.traces = np.array(padded_traces)
            else:
                raise ValueError("No power trace data found in NPZ file")
            
            # Optional字段
            self.timestamps = data.get('timestamps')
            self.durations = data.get('durations')
            
            self.logger.info(f"Data loaded successfully:")
            self.logger.info(f"  - Traces: {self.traces.shape}")
            self.logger.info(f"  - Plaintexts: {self.plaintexts.shape}")
            if self.keys is not None:
                self.logger.info(f"  - Keys: {self.keys.shape}")
            else:
                self.logger.info(f"  - Keys: None (will be recovered by attack)")
            
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to load data: {e}")
            return False
    
    def convert_bytes_to_int(self, byte_data: np.ndarray) -> np.ndarray:
        """将字节ArrayConvertForIntegerArray"""
        if byte_data.dtype.kind in ['S', 'U']:  # StringClass型
            return np.array([pt[0] if len(pt) > 0 else 0 for pt in byte_data])
        elif byte_data.dtype == object:  # 对象Array（字节Data）
            # Process字节Data，ConvertForIntegerMatrix
            result = []
            for item in byte_data:
                if isinstance(item, bytes):
                    result.append(list(item))
                elif isinstance(item, (list, np.ndarray)):
                    result.append(list(item))
                else:
                    result.append([int(item)])
            return np.array(result)
        else:
            return byte_data
    
    @abstractmethod
    def attack(self, target_byte: int = 0, **kwargs) -> AttackResult:
        """
        ExecuteAttack
        
        Args:
            target_byte: 目标字节位置
            **kwargs: 其他AttackParameters
            
        Returns:
            AttackResult
        """
        pass
    
    @abstractmethod
    def plot_results(self, result: AttackResult) -> None:
        """
        绘制AttackResult
        
        Args:
            result: AttackResult
        """
        pass
    
    def validate_data(self) -> bool:
        """VerifyDataYesNo已Load且Has效"""
        if self.traces is None:
            self.logger.error("No trace data loaded")
            return False
        
        if self.plaintexts is None:
            self.logger.error("No plaintext data loaded")
            return False
            
        if len(self.traces) != len(self.plaintexts):
            self.logger.error("Trace and plaintext data size mismatch")
            return False
            
        return True
    
    def get_target_data(self, target_byte: int) -> np.ndarray:
        """
        获取目标字节Data
        
        Args:
            target_byte: 目标字节位置
            
        Returns:
            目标字节DataArray
        """
        plaintext_ints = self.convert_bytes_to_int(self.plaintexts)
        
        if len(plaintext_ints.shape) > 1:
            if target_byte >= plaintext_ints.shape[1]:
                raise ValueError(f"Target byte {target_byte} exceeds data width {plaintext_ints.shape[1]}")
            return plaintext_ints[:, target_byte]
        else:
            if target_byte > 0:
                raise ValueError(f"Target byte {target_byte} not available in single-byte data")
            return plaintext_ints
    
    def calculate_intermediate_values(self, target_data: np.ndarray, hypothesis: int) -> np.ndarray:
        """
        CalculateIntermediate value（可被子Class重写以实现特定的Intermediate valueCalculate）
        
        Args:
            target_data: 目标Data
            hypothesis: False设值
            
        Returns:
            Intermediate valueArray
        """
        # DefaultUseSimple的XOR操作
        return np.bitwise_xor(target_data, hypothesis)
    
    def save_results(self, result: AttackResult, output_dir: str = "attack_results") -> str:
        """
        SaveAttackResult
        
        Args:
            result: AttackResult
            output_dir: Output directory
            
        Returns:
            Save的FilePath
        """
        import os
        from datetime import datetime
        
        os.makedirs(output_dir, exist_ok=True)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{self.name}_{timestamp}.npz"
        filepath = os.path.join(output_dir, filename)
        
        # SaveResult
        np.savez_compressed(
            filepath,
            attack_type=result.attack_type,
            target_byte=result.target_byte,
            execution_time=result.execution_time,
            success=result.success,
            confidence=result.confidence,
            best_hypothesis=result.best_hypothesis,
            metadata=result.metadata,
            timestamp=timestamp
        )
        
        self.logger.info(f"Results saved to: {filepath}")
        return filepath


class PowerAnalysisAttack(BaseSideChannelAttack):
    """Power AnalysisAttack基Class"""
    
    def __init__(self, name: str):
        super().__init__(name)
        self.hypothesis_range = 256  # DefaultFalse设范围（8位）
    
    def set_hypothesis_range(self, range_size: int):
        """SettingsFalse设范围"""
        self.hypothesis_range = range_size
    
    def calculate_correlation(self, intermediate_values: np.ndarray, sample_point: int) -> float:
        """
        Calculate相关性系数（带数值稳定性Process）
        
        Args:
            intermediate_values: Intermediate valueArray
            sample_point: Sampling点位置
            
        Returns:
            相关性系数
        """
        if sample_point >= self.traces.shape[1]:
            return 0.0
        
        # 获取PowerTraceData
        power_values = self.traces[:, sample_point]
        
        # CheckDataHas效性
        if len(intermediate_values) != len(power_values):
            return 0.0
        
        # Check标准差，Avoid除零
        std_intermediate = np.std(intermediate_values)
        std_power = np.std(power_values)
        
        if std_intermediate == 0.0 or std_power == 0.0:
            return 0.0
        
        # Use更稳定的相关性Calculate
        try:
            # Add小的噪声以Avoid数值问题
            epsilon = 1e-10
            intermediate_stable = intermediate_values + np.random.normal(0, epsilon, len(intermediate_values))
            power_stable = power_values + np.random.normal(0, epsilon, len(power_values))
            
            # Calculate相关性
            with np.errstate(divide='ignore', invalid='ignore'):
                corr = np.corrcoef(intermediate_stable, power_stable)[0, 1]
            
            return 0.0 if np.isnan(corr) or np.isinf(corr) else corr
            
        except Exception:
            return 0.0


class StatisticalAttack(BaseSideChannelAttack):
    """统计AnalyzeAttack基Class"""
    
    def __init__(self, name: str):
        super().__init__(name)
        self.significance_threshold = 0.05  # Default显著性阈值
    
    def set_significance_threshold(self, threshold: float):
        """Settings显著性阈值"""
        self.significance_threshold = threshold
    
    def perform_t_test(self, group1: np.ndarray, group2: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Executet检验
        
        Args:
            group1: First组Data
            group2: 第二组Data
            
        Returns:
            t统计量Andp值Array
        """
        from scipy import stats
        return stats.ttest_ind(group1, group2, axis=0)


class TemplateBasedAttack(BaseSideChannelAttack):
    """Based on模板的Attack基Class"""
    
    def __init__(self, name: str):
        super().__init__(name)
        self.templates = []
        self.num_templates = 5
    
    def set_num_templates(self, num_templates: int):
        """Settings模板数量"""
        self.num_templates = num_templates
    
    def create_templates(self) -> List[Dict[str, Any]]:
        """
        Create模板
        
        Returns:
            模板List
        """
        from sklearn.cluster import KMeans
        
        # UseK-means聚ClassCreate模板
        kmeans = KMeans(n_clusters=self.num_templates, random_state=42)
        traces_reshaped = self.traces.reshape(self.traces.shape[0], -1)
        cluster_labels = kmeans.fit_predict(traces_reshaped)
        
        templates = []
        for i in range(self.num_templates):
            cluster_traces = self.traces[cluster_labels == i]
            if len(cluster_traces) > 0:
                template = {
                    'cluster_id': i,
                    'mean': np.mean(cluster_traces, axis=0),
                    'covariance': np.cov(cluster_traces.T),
                    'size': len(cluster_traces)
                }
                templates.append(template)
        
        self.templates = templates
        return templates
